import matplotlib.pyplot as plt
import numpy as np

# 设置随机数种子
np.random.seed(42)

# %matplotlib inline

# 集群数
K = 3
# 两个维度
D = 2
# 点的总数
N = 1000
# 不同集群的点数
Ns = [300, 400, 300]

# 3行2列 的随机数组(服从正态分布)
# 相当于生成3个二维坐标点 -> [[x1,y1],[x2,y2],[x3,y3]]
means = 5 * np.random.randn(K, D)
print(means)

x = []
# x: ['a','b','c'], y: [1, 2, 3] -> zip(x,y) -> ('a',1), ('b', 2), ('c',3)
# zip(Ns, means) -> [(300,[x1,y1]),(400,[x2,y2]),(300,[x3,y3])]
for n, m in zip(Ns, means):
    print(n, m)
    # 生成 n行2列的 随机数组 -> 相当于n个二维坐标点 -> m是means里的二维坐标(集群中心点)
    x.append(np.random.randn(n, D) + m)

# 此时x的shape(形状)
# [(300, 2), (400, 2), (300, 2)]
print([x_.shape for x_ in x])

# zip -> [((300, 2),[x1,y1]),((400, 2),[x2,y2]),((300, 2),[x3,y3])]
for x_, m in zip(x, means):
    print(x_.shape, m.shape)
    # scatter(x,y) -> 绘制散点图
    # x: n行2列 -> x_[:,0]: 所有行的第一列(xi), x[:,1]: 所有行的第二列(yi)
    plt.scatter(x_[:, 0], x_[:, 1])
    # 绘制集群中心点
    # plt.plot里的这个'kx' -> 颜色字符k: 表示黑色，标记字符x: 表示图案为x标记
    plt.plot(m[0], m[1], 'kx')
plt.title('True Clusters')
plt.show()

# 从x取出所有数据
# np.vstack: 按垂直方向（行顺序）堆叠数组构成一个新的数组,堆叠的数组需要具有相同的维度
# np.hstack: 按水平方向（列顺序）堆叠数组构成一个新的数组,堆叠的数组需要具有相同的维度
data = np.vstack(x)
print(data.shape)
plt.scatter(data[:, 0], data[:, 1])
plt.show()


# 计算欧式距离 √((x1-x2)^2 + (y1-y2)^2)
def dist(x1, x2, axis=None):
    """
    Calculate euclidean distance
    :param x1: 点1 (x1,y1)
    :param x2: 点2 (x2,y2)
    :param axis:
    :return:
    """
    # np.square(x1 - x2) -> (x1-x2)^2 平方
    # np.sqrt 开根号
    return np.sqrt(np.sum(np.square(x1 - x2), axis))


# 计算平均值和任意给定点的距离
def distance_matrix(x, m):
    """
    Calculates the distance from each element x to each element in m.
    :param x: data points
    :param m: possible means
    :return: distance matrix
    """
    # 初始化(用0填充)距离矩阵(len(x)行len(m)列)
    d = np.zeros((len(x), len(m)))
    for i in range(len(x)):
        for j in range(len(m)):
            # 传入x的第i行所有列 x[i] -> (xi,yi)
            # 传入m的第j行所有列 m[j] -> (xm,ym)
            d[i, j] = dist(x[i, :], m[j, :])

    return d


# 测试
# Test out the distance matrix algorithm
x_test = np.array([[0, 1], [1, 0], [0, 0], [1, 1]])
print(distance_matrix(x_test, x_test))


# 使用广播替代循环(python的循环性能低)
# broadcasting
# 计算欧式距离 √((x1-x2)^2 + (y1-y2)^2)
def dist(x1, x2, axis=-1):
    """
    Calculate euclidean distance
    :param x1: 点1 (x1,y1)
    :param x2: 点2 (x2,y2)
    :param axis:
    :return:
    """
    return np.sqrt(np.sum(np.square(x1 - x2), axis))


def distance_matrix_broadcasting(x, m):
    # np.sum(axis=0) -> 行(维度1)求和 -> [[x1,y1],[x2,y2]] -> [[x1+x2],[y1+y2]]
    # np.sum(axis=1) -> 列(维度2)求和 -> [[x1,y1],[x2,y2]] -> [[x1+y1],[x2+y2]]
    # x所有点与m_相减([xx-mx,xy-my])并求平方，然后列相加((xx-mx)^2+(xy-my)^2，然后开根号得距离
    d = [dist(x, m_, axis=1) for m_ in m]
    # print(d)
    # 按行堆叠(axis=1)
    d = np.stack(d, axis=1)
    return d


# 测试
# Test out the distance matrix algorithm
x_test = np.array([[0, 1], [1, 0], [0, 0], [1, 1]])
print(distance_matrix_broadcasting(x_test, x_test))

# 计算k-means

# 初始化数据
# 集群数
k = 3
#
iters = 10
print(data.shape)
print(data)
# (1000, 2)
# 3个集群中心点，(3行2列)
means = np.random.randn(k, data.shape[1])

for i in range(iters):
    # 求data和中心点的距离 -> shape: (1000, 3) -> 这1000个点分别到3个中心点的距离
    d = distance_matrix_broadcasting(data, means)
    # print(d[:10])
    # print(d.shape) # (1000,3)
    # np.argmin: 找到最小的(值的下标)然后映射到指定轴
    '''
    array([[0, 4, 2],
           [3, 1, 5]]) # 3*2
    >>> np.argmin(a, axis=0)
    映射到行x
    找出每列最小元素的下标
    array([0 1 0]) # (3,)
    >>> np.argmin(a, axis=1)
    映射到列y
    找出每行最小元素的下标
    array([0 1]) #(2,)
    '''
    cluster = d.argmin(axis=-1)  # 找到平均距离最小值然后映射到列
    # true
    # print(cluster == np.argmin(d, axis=1))
    # print(cluster[:10])  # [一堆0、1、2]
    # print(cluster.shape) # (1000,)
    for j in range(k):
        # print(j) # 0,1,2
        # 筛选出和第j个集群中心点最近的下标
        idx = cluster == j  # numpy布尔索引
        # print(idx)
        # print(data[idx])
        # 绘制集群中心点
        plt.plot(means[j, 0], means[j, 1], 'rx')
        # data[idx] -> 找出所有cluster==j(与集群中心点j最近)的元素，组成新数组
        plt.scatter(data[idx, 0], data[idx, 1])
        # calculate new mean
        # print('means\n', means[j, :])
        # print('new:\n', data[idx].mean(axis=0))
        # 根据当前平均距离最小的点确定新点
        # 将距离集群中心点j最近的所有点的坐标求平均,得出新点的坐标
        means[j, :] = data[idx].mean(axis=0)
        # print(data[idx])
        # print(means[j, :])
    plt.show()
